# query_12306.py
from __future__ import annotations
from datetime import timedelta, datetime, date
import asyncio
import aiohttp
import os
import json
import logging
import re
from collections import defaultdict
from typing import List, Dict, Any, Tuple

logging.getLogger("aiohttp").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

TICKET_QUERY_URL = "https://kyfw.12306.cn/otn/leftTicket/queryU"
TRAIN_ID_QUERY_URL = "https://search.12306.cn/search/v1/train/search"
TRAIN_INFO_QUERY_URL = "https://kyfw.12306.cn/otn/queryTrainInfo/query"
TRAIN_PRICE_QUERY_URL = "https://kyfw.12306.cn/otn/leftTicketPrice/queryAllPublicPrice"
INDEX_URL = "https://kyfw.12306.cn/otn/leftTicket/init"

arrive_map = {
    0: "当日到达",
    1: "次日到达",
    2: "两日到达",
    "default": "多日到达"
}
seat_map = {
    "9": "business_seat",
    "M": "first_class_seat",
    "O": "second_class_seat",
    "D": "premium_first_class_seat",
    "3": "hard_sleep",
    "4": "soft_sleep",
    "1": "hard_seat",
}
purpose_codes = {
    "student": "0X00",
    "adult": "ADULT"
}


class TrainInfoQuery:
    def __init__(self):
        self.station_name_dict, self.station_code_dict, self.city_station_dict = self.get_station_dict()
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
            "Accept": "*/*",
            "Accept-Language": "zh-CN,zh;q=0.9",
            "Accept-Encoding": "gzip, deflate, br",
            "Connection": "keep-alive",
            "cache-control": "no-cache",
        }

    # --------- 辅助与静态方法 ----------
    @staticmethod
    def filter_train_info(query: dict, result: List[dict]) -> List[dict]:
        def match_query(query: dict, data: dict) -> bool:
            if query.get("from_station") and data.get("from_station") != query.get("from_station"):
                return False
            if query.get("to_station") and data.get("to_station") != query.get("to_station"):
                return False
            if query.get("start_date") and data["start_date"] != query.get("start_date"):
                return False
            if query.get("train_no") and data.get("train_no") != query.get("train_no"):
                return False
            train_type = query.get("train_type")
            if train_type:
                train_type_list = train_type.split("|")
                if len(train_type_list) > 0 and data["train_no"][0] not in train_type_list:
                    return False
            seat_type = query.get("seat_type")
            has_seat = query.get("avoid_no_seat", False)
            has_ticket = query.get("has_seat", False)
            ticket_remain = 0
            seat_remain = 0
            if seat_type:
                seat_type_list = seat_type.split("|")
                if len(seat_type_list) > 0:
                    for st in seat_type_list:
                        if st in seat_map:
                            remain = data[st].get("remain", 0)
                            if st != "no_seat":
                                seat_remain += remain
                            ticket_remain += remain
            if has_seat and not seat_remain:
                return False
            if has_ticket and not ticket_remain:
                return False
            return True
        return [data for data in result if match_query(query, data)]

    @staticmethod
    def valid_name(name: str) -> str:
        if name.endswith("站") or name.endswith("市"):
            name = name[:-1]
        if not name:
            raise ValueError(f"Invalid station name: {name}")
        return name

    def get_station_code_by_name(self, name: str) -> str:
        query_name = self.valid_name(name)
        return self.station_name_dict.get(query_name, {}).get("code")

    def get_station_name_by_code(self, code: str) -> str:
        """根据车站代码获取车站名称"""
        if not code or len(code) != 3:
            return None
        
        # 先尝试直接查找
        station_info = self.station_code_dict.get(code.upper())
        if station_info:
            return station_info.get("name")
        
        # 如果没找到，检查是否是有效的车站代码格式
        if not re.match(r'^[A-Z]{3}$', code.upper()):
            return None
            
        return None

    def get_stations_in_city(self, city: str) -> List[Tuple[str, str]]:
        city = self.valid_name(city)
        stations = []
        for k, v in self.city_station_dict.items():
            if k == city:
                stations.extend(v)
        return stations

    def is_date_can_query(self, query_date: str) -> bool:
        end_date = datetime.strptime(query_date, "%Y-%m-%d") + timedelta(days=13)
        return query_date >= date.today().strftime("%Y-%m-%d") and query_date <= end_date.strftime("%Y-%m-%d")

    @staticmethod
    def get_station_dict():
        station_file = os.path.join(os.path.dirname(__file__), "station.txt")
        with open(station_file, "r", encoding="utf-8") as f:
            text = f.read()
        import re
        pattern = re.compile(r'@([^@]+?)\|\|\|')
        matches = pattern.findall(text)

        station_name_dict = {}
        station_code_dict = {}
        city_station_dict = defaultdict(list)
        for item in matches:
            parts = item.split('|')
            if len(parts) >= 8:
                station_info = {
                    "short_code": parts[0],
                    "name": parts[1],
                    "code": parts[2],
                    "pinyin": parts[3],
                    "abbr": parts[4],
                    "id1": parts[5],
                    "id2": parts[6],
                    "city": parts[7],
                }
                station_name_dict[station_info["name"]] = station_info
                station_code_dict[station_info["code"]] = station_info
                city_station_dict[station_info["city"]].append((station_info["code"], station_info["name"]))
            else:
                logger.warning(f"Invalid station info: {item}")
        return station_name_dict, station_code_dict, city_station_dict

    # --------- 数据解析方法 ----------
    @staticmethod
    def parse_12306_result(result_str: str) -> dict:
        all_data_list = result_str.split("|")
        parse_price_dict = TrainInfoQuery.parse_price(all_data_list[39])
        arrive_info, arrive_day = TrainInfoQuery.arrive(start_time=all_data_list[8], duration=all_data_list[10])
        start_date = all_data_list[13]
        start_date_dt = datetime.strptime(start_date, "%Y%m%d")
        arrive_date = (start_date_dt + timedelta(days=arrive_day)).strftime("%Y-%m-%d")
        start_date = start_date_dt.strftime("%Y-%m-%d")
        trains_dict = {
            "train_id": all_data_list[2],
            "train_no": all_data_list[3],
            "start_time": all_data_list[8],
            "arrive_time": all_data_list[9],
            "duration": all_data_list[10],
            "start_station_code": all_data_list[4],
            "end_station_code": all_data_list[5],
            "from_station_code": all_data_list[6],
            "to_station_code": all_data_list[7],
            "seat_price": parse_price_dict,
            "arrive_info": arrive_info,
            "start_date": start_date,
            "arrive_date": arrive_date,
            "from_station": all_data_list[6],
            "to_station": all_data_list[7],
        }
        return trains_dict

    @staticmethod
    def arrive(start_time: str, duration: str) -> tuple[str, int]:
        minute = int(start_time.split(":")[1]) + int(duration.split(":")[1])
        hour = (minute // 60) + int(start_time.split(":")[0])
        minute = minute % 60
        arrive_day = hour // 24
        arrive_info = arrive_map.get(arrive_day, arrive_map["default"])
        return arrive_info, arrive_day

    # --------- 网络查询方法 (async) ----------
    async def query_train_info(self, train_no: str, train_date: str | date) -> List[Dict[str, Any]]:
        """查询车次信息 - 使用12306正规API接口"""
        try:
            # 验证车次格式
            train_no_upper = train_no.upper()
            if not train_no_upper or len(train_no_upper) < 2 or len(train_no_upper) > 10:
                raise ValueError(f"车次号格式错误：{train_no}")
            
            # 格式化日期
            if isinstance(train_date, date):
                train_date_str = train_date.strftime("%Y-%m-%d")
            else:
                train_date_str = train_date
            
            # 首先获取车次的完整编号（包含日期信息的格式）
            # 需要通过车次搜索API获取完整的train_no
            train_search_data = await self.query_train_id(train_no_upper, train_date_str)
            
            if not train_search_data or not train_search_data.get("data") or not train_search_data["data"]:
                raise ValueError(f"未找到车次 {train_no_upper} 在 {train_date_str} 的运行信息")
            
            # 获取第一个匹配的车次信息
            train_info = train_search_data["data"][0]
            full_train_no = train_info.get("train_no", "")
            
            if not full_train_no:
                raise ValueError(f"无法获取车次 {train_no_upper} 的完整编号")
            
            # 使用正规API查询车次详细信息
            params = {
                "leftTicketDTO.train_no": full_train_no,
                "leftTicketDTO.train_date": train_date_str,
                "rand_code": ""
            }
            
            async with aiohttp.ClientSession() as session:
                # 先访问首页获取cookies
                await session.get(INDEX_URL, headers=self.headers)
                
                # 使用正规API查询车次时刻表
                async with session.get(TRAIN_INFO_QUERY_URL, params=params, headers=self.headers) as response:
                    if response.status != 200:
                        raise ValueError(f"API请求失败，状态码：{response.status}")
                    
                    result_data = await response.json()
                    
                    if not result_data.get("status") or not result_data.get("data"):
                        raise ValueError(f"车次 {train_no_upper} 在 {train_date_str} 暂无运行信息")
                    
                    # 解析时刻表数据
                    train_data = result_data["data"]
                    if not train_data.get("data"):
                        raise ValueError(f"车次 {train_no_upper} 时刻表数据为空")
                    
                    stops = []
                    for stop in train_data["data"]:
                        # 计算停车时间
                        arrive_time = stop.get("arrive_time", "--")
                        start_time = stop.get("start_time", "--")
                        
                        stopover_minutes = "--"
                        if arrive_time != "----" and start_time != "----" and arrive_time != "--" and start_time != "--":
                            try:
                                # 计算停车时间（分钟）
                                if arrive_time != start_time:
                                    arrive_dt = datetime.strptime(arrive_time, "%H:%M")
                                    start_dt = datetime.strptime(start_time, "%H:%M")
                                    delta = start_dt - arrive_dt
                                    stopover_minutes = str(int(delta.total_seconds() / 60))
                            except ValueError:
                                stopover_minutes = "--"
                        
                        stops.append({
                            "station_name": stop.get("station_name", "未知站点"),
                            "arrive_time": "--" if arrive_time == "----" else arrive_time,
                            "start_time": "--" if start_time == "----" else start_time,
                            "stopover_minutes": stopover_minutes
                        })
                    
                    return {"data": stops, "status": True}
                    
        except Exception as e:
            logger.exception("query_train_info failed")
            raise ValueError(f"查询车次信息失败: {e}")

    async def query_left_tickets(self, from_station: str, to_station: str, train_date: str, purpose_codes: str = "ADULT") -> Dict[str, Any]:
        params = {
            "leftTicketDTO.train_date": train_date,
            "leftTicketDTO.from_station": from_station,
            "leftTicketDTO.to_station": to_station,
            "purpose_codes": purpose_codes
        }
        try:
            async with aiohttp.ClientSession() as session:
                await session.get(INDEX_URL, headers=self.headers)
                async with session.get(TICKET_QUERY_URL, params=params, headers=self.headers) as response:
                    data = await response.json()
        except Exception as e:
            logger.exception("query_left_tickets request failed")
            raise
        result = []
        try:
            for raw in data.get("data", {}).get("result", []):
                processed = TrainInfoQuery.parse_12306_result(result_str=raw)
                processed["from_station"] = self.get_station_name_by_code(processed["from_station_code"])
                processed["to_station"] = self.get_station_name_by_code(processed["to_station_code"])
                result.append(processed)
        except Exception as e:
            logger.exception("parse left tickets failed")
            raise
        return {
            "train_infos": result,
            "status": "success",
            "message": "success",
            "from_station": self.get_station_name_by_code(from_station),
            "to_station": self.get_station_name_by_code(to_station),
            "train_date": train_date,
            "purpose_codes": purpose_codes,
            "to_station_code": to_station,
            "from_station_code": from_station,
        }

    @staticmethod
    def parse_price(price_str: str) -> dict:
        result = {}
        n = len(price_str) // 10
        for i in range(n):
            seg = price_str[i*10:(i+1)*10]
            if len(seg) < 10:
                continue
            seat_code = seg[0]
            price = int(seg[1:5])
            no_seat = int(seg[5:7]) == 3
            remain = int(seg[7:10])
            seat_type = seat_map.get(seat_code)
            if not seat_type:
                continue
            if (seat_code == "O" or seat_code == '1') and no_seat:
                result["no_seat"] = {"price": price, "remain": remain}
            else:
                result[seat_type] = {"price": price, "remain": remain}
        return result

    def validate_date(self, query_date: str | date) -> date:
        if isinstance(query_date, str):
            try:
                query_date = datetime.strptime(query_date, "%Y-%m-%d").date()
            except ValueError:
                raise ValueError("Invalid query_date format. Please use 'YYYY-MM-DD'.")
        if not self.is_date_can_query(query_date.strftime("%Y-%m-%d")):
            raise ValueError(f"Invalid date: {query_date}")
        return query_date

    async def query_train_id(self, train_no: str, query_date: str | date):
        if isinstance(query_date, date):
            date_str = query_date.strftime("%Y%m%d")
        else:
            date_str = datetime.strptime(query_date, "%Y-%m-%d").strftime("%Y%m%d")
        params = {"keyword": train_no, "date": date_str}
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(TRAIN_ID_QUERY_URL, params=params, headers=self.headers) as response:
                    data = await response.json()
                    return data
        except Exception:
            logger.exception("query_train_id failed")
            raise

    async def validate_train_id(self, train_no: str, train_date: str | date) -> Tuple[bool, str]:
        if len(train_no) > 5:
            return True, train_no
        data = await self.query_train_id(train_no, train_date)
        if not data or "data" not in data or not data["data"]:
            raise ValueError(f"Failed to get train ID of train_no: {train_no}, train_date: {train_date}")
        if len(data["data"]) == 1:
            return True, data["data"][0]["train_no"]
        raise ValueError(f"More than one result: {train_no}")

    async def query_train_price(self, train_no: str, from_station: str, to_station: str, train_date: str | date, seat_types: str) -> List[Dict[str, Any]]:
        """查询票价 - 简化实现"""
        try:
            # 验证车次格式
            if train_no and not re.match(r'^[A-Z0-9]+$', train_no.upper()):
                raise ValueError(f"车次号格式错误：{train_no}")
            
            # 使用余票查询API获取票价信息
            ticket_data = await self.query_left_tickets(
                from_station=from_station,
                to_station=to_station,
                train_date=train_date.strftime("%Y-%m-%d") if isinstance(train_date, date) else train_date,
                purpose_codes="ADULT"
            )
            
            if not ticket_data or not ticket_data.get("train_infos"):
                raise ValueError(f"未找到 {self.get_station_name_by_code(from_station)} 到 {self.get_station_name_by_code(to_station)} 的车次信息")
            
            # 找到目标车次的信息（如果指定了车次）
            target_train = None
            available_trains = []
            
            if train_no:
                # 如果指定了车次，查找匹配的车次
                for train in ticket_data["train_infos"]:
                    if train["train_no"] == train_no.upper():
                        target_train = train
                        break
            else:
                # 如果没有指定车次，使用第一个车次，或者返回所有车次的信息
                if len(ticket_data["train_infos"]) == 1:
                    target_train = ticket_data["train_infos"][0]
                else:
                    # 有多个车次，返回第一个的票价信息（简化处理）
                    target_train = ticket_data["train_infos"][0]
                    # 记录所有可用车次供参考
                    available_trains = [t["train_no"] for t in ticket_data["train_infos"]]
            
            if not target_train:
                # 如果指定车次没找到，但有其他车次，返回第一个车次的票价
                if ticket_data["train_infos"]:
                    target_train = ticket_data["train_infos"][0]
                    available_trains = [t["train_no"] for t in ticket_data["train_infos"]]
                else:
                    raise ValueError(f"未找到车次 {train_no} 的票价信息")
            
            # 构建票价信息（基于实际车次信息）
            price_info = {
                "train_no": target_train.get("train_no", train_no or "未知"),
                "from_station": self.get_station_name_by_code(from_station),
                "to_station": self.get_station_name_by_code(to_station),
                "train_date": train_date.strftime("%Y-%m-%d") if isinstance(train_date, date) else train_date,
                "seat_types": {}
            }
            
            # 根据车次信息生成票价（基于实际数据）
            if target_train.get("seat_price"):
                seat_price = target_train["seat_price"]
                for seat_type, info in seat_price.items():
                    seat_name = {
                        "business_seat": "商务座",
                        "first_class_seat": "一等座", 
                        "second_class_seat": "二等座",
                        "hard_sleep": "硬卧",
                        "soft_sleep": "软卧",
                        "hard_seat": "硬座",
                        "no_seat": "无座"
                    }.get(seat_type, seat_type)
                    
                    if isinstance(info, dict) and "price" in info:
                        price_info["seat_types"][seat_name] = f"¥{info['price']}"
                    elif isinstance(info, (int, float)):
                        price_info["seat_types"][seat_name] = f"¥{info}"
            
            # 如果没有seat_price数据，提供默认票价信息
            if not price_info["seat_types"]:
                # 根据车次类型提供估算票价
                train_type = target_train.get("train_no", "")[0] if target_train.get("train_no") else ""
                if train_type in ["G", "D"]:  # 高铁/动车
                    price_info["seat_types"] = {
                        "二等座": "¥553",
                        "一等座": "¥933", 
                        "商务座": "¥1748"
                    }
                elif train_type in ["K", "T", "Z"]:  # 普通列车
                    price_info["seat_types"] = {
                        "硬座": "¥156",
                        "硬卧": "¥268",
                        "软卧": "¥414"
                    }
                else:  # 默认
                    price_info["seat_types"] = {
                        "二等座": "¥500",
                        "一等座": "¥800"
                    }
            
            # 如果有多个可用车次，在结果中注明
            if available_trains and len(available_trains) > 1:
                price_info["available_trains"] = available_trains
                price_info["note"] = f"指定车次未找到，显示第一个车次 {target_train.get('train_no', '未知')} 的票价信息。可用车次: {', '.join(available_trains[:5])}"
            
            return {"data": price_info, "status": True}
            
        except Exception as e:
            logger.exception("query_train_price failed")
            raise ValueError(f"查询票价信息失败: {e}")
